package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo4TrainModel {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[7]")
      .appName("image")
      .config("spark.sql.shuffle.partitions", 8)
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._

    //读取处理好的图片向量数据
    val imageDF: DataFrame = spark
      .read
      .format("libsvm")
      .option("numFeatures", 784)
      .load("data/image_data")

    /**
     * 将原始数据拆分成训练集和测试集
     * 训练集：用于训练模型
     * 测试集：用于测试模型准确率
     */
    val Array(train, test) = imageDF.randomSplit(Array(0.7, 0.3))

    /**
     * 选择算法，训练模型
     */

    val lr = new LogisticRegression()

    //将训练姐带入算法训练模型
    val mode: LogisticRegressionModel = lr.fit(train)

    /**
     * 将测试集带入模型测试模型的准确率
     * 准确率：对的数量/总的数量
     * 精确率：即正确预测为正的占全部预测为正的比例
     * 召回率：即正确预测为正的占全部实际为正的比例
     */


    val frame: DataFrame = mode.transform(test)

    //计算准确率
    val acc: Double = frame.where($"label" === $"prediction").count().toDouble / frame.count()

    println(s"准确率：$acc")

    /**
     * 保存模型
     */
    mode
      .write
      .overwrite()
      .save("data/image_model")

  }

}
